import os
import pandas as pd
import xlsxwriter as xl
from xlsxwriter.utility import xl_range
from utils import *

def load_professional_score_data():
    """加载专业分数数据"""
    csv_path = './data/professional_score.csv'
    if not os.path.exists(csv_path):
        print(f"专业分数数据文件不存在: {csv_path}")
        return None
    return pd.read_csv(csv_path)

def build_excel_table(science_rank_df, arts_rank_df, title, file):
    """构建Excel表格"""
    os.makedirs(os.path.dirname(file), exist_ok=True)
    writer = pd.ExcelWriter(file, engine='xlsxwriter')
    
    subjects_data = {
        'science': {'data': science_rank_df, 'name': '理工类'},
        'arts': {'data': arts_rank_df, 'name': '文史类'}
    }
    
    for subject, info in subjects_data.items():
        if info['data'].empty:
            continue
            
        # 准备数据，只选择排名列
        table = info['data'].copy()
        
        # 重新排列列顺序，将排名列按年份从新到旧排序，移除院校代码列
        base_cols = ['school_name', 'spname', 'zs_type', 'sg_name', 'local_batch_name']
        rank_cols = [f'rank_{year}' for year in range(2024, 2017, -1)]
        
        # 确保所有列都存在
        for col in base_cols + rank_cols:
            if col not in table.columns:
                table[col] = None
                
        table = table[base_cols + rank_cols]
        
        # 重命名列为中文，移除院校代码映射
        column_mapping = {
            'school_name': '院校名称', 
            'spname': '专业名称',
            'zs_type': '招生类型',
            'sg_name': '专业组',
            'local_batch_name': '批次'
        }
        
        for year in range(2024, 2017, -1):
            column_mapping[f'rank_{year}'] = f'名次{year}'
            
        table.rename(columns=column_mapping, inplace=True)
        
        # 按最新年份名次排序
        latest_rank_col = f'名次{2024}'
        if latest_rank_col in table.columns:
            table = table.sort_values(by=latest_rank_col, na_position='last')
        
        sheet_name = f'历史名次（{info["name"]}）'
        _, _, nrow, ncol = write_to_excel_table(table, writer, sheet_name, title + f'（{info["name"]}）')
        workbook = writer.book
        worksheet = writer.sheets[sheet_name]

        # 添加说明
        red = workbook.add_format({'color': '#C00000', 'bold': 1})
        blue = workbook.add_format({'color': '#16365C', 'bold': 1})
        cell_format = workbook.add_format({'align': 'left',
                                    'valign': 'top',
                                    'text_wrap': True})
        
        INFO_ROW_START = 2
        INFO_ROW_END = 18
        IMAGE_ROW_START = 19
        IMAGE_ROW_END = IMAGE_ROW_START + 12
        worksheet.merge_range(INFO_ROW_START, ncol+1, INFO_ROW_END, ncol+1, '', cell_format)
        worksheet.set_column_pixels(ncol+1, ncol+1, 220)
        worksheet.write_rich_string(2, ncol+1, 
            blue, '数据说明：\n',
            red, '数据较多，请善用 Excel 筛选功能。\n',
            '数据采用历年专业录取最低排名。空白代表当年该专业数据缺失。\n',
            '某些排名要求较高的专业（如清北的很多专业）没有具体排名数据，因此不显示在表内。\n',
            red, '仅按专业名称判断专业相同', '，因此可能存在同名但实际不同的专业被错误归类到一起。\n',
            red, '数据按个人理解解析，首次制作，可能存在大量错误，无法保证准确！！！！\n',
            red, '因此，数据仅供参考！！！！请以考试院官方发布数据为准。', cell_format
        )
        cell_format = workbook.add_format({'align': 'center',
                                    'valign': 'top',
                                    'color': '#16365C',
                                    'bold': 1,
                                    'text_wrap': True})
        worksheet.merge_range(IMAGE_ROW_START, ncol+1, IMAGE_ROW_END, ncol+1, '欢迎报考\n中国科学技术大学', cell_format)
        if 'PATH_USTC_LOGO' in globals():
            worksheet.insert_image(IMAGE_ROW_START+2, ncol+1, PATH_USTC_LOGO, {'x_scale': 1, 'y_scale': 5 / 4.52, 'x_offset': 30})

    writer.close()

def build_subject_rank_df(subject_data, subject_name):
    """为指定科目构建历史排名DataFrame"""
    print(f"\n正在处理{subject_name}数据...")
    
    # 只保留批次名称中包含"本科"的记录
    subject_data = subject_data[subject_data['local_batch_name'].str.contains('本科', na=False)].copy()
    print(f"  过滤后剩余 {len(subject_data)} 条本科专业记录")
    
    # 按学校、专业和招生类型分组
    grouped = subject_data.groupby(['school_id', 'school_name', 'spname', 'zs_type'])
    
    rank_records = []
    total_groups = len(grouped)
    processed_groups = 0
    
    for (school_id, school_name, spname, zs_type), group in grouped:
        processed_groups += 1
        if processed_groups % 100 == 0:
            print(f"  处理进度: {processed_groups}/{total_groups}")
        
        # 获取最新年份的专业组和批次信息
        latest_year = group['year'].max()
        latest_year_data = group[group['year'] == latest_year].iloc[-1]
        sg_name = latest_year_data.get('sg_name', '')
        local_batch_name = latest_year_data.get('local_batch_name', '')
        
        # 构建历年排名数据，直接使用min_section作为排名
        year_ranks = {}
        
        for _, row in group.iterrows():
            year = row['year']
            min_section = row['min_section']
            
            # 检查min_section是否为有效数值（不是NaN、None或"-"）
            if pd.notna(min_section) and str(min_section).strip() != "-" and str(min_section).strip() != "":
                try:
                    year_ranks[f'rank_{year}'] = int(float(min_section))
                except (ValueError, TypeError):
                    # 如果无法转换为数字，跳过这个数据
                    continue
        
        # 只有当至少有一个年份有排名数据时才添加记录
        if year_ranks:
            record = {
                'school_id': school_id,
                'school_name': school_name,
                'spname': spname,
                'zs_type': zs_type,
                'sg_name': sg_name,
                'local_batch_name': local_batch_name,
                **year_ranks
            }
            rank_records.append(record)
    
    print(f"  {subject_name}完成处理，共 {len(rank_records)} 个有效记录")
    
    if rank_records:
        df = pd.DataFrame(rank_records)
        
        # 确保所有年份的排名列都存在
        for year in range(2018, 2025):
            col_name = f'rank_{year}'
            if col_name not in df.columns:
                df[col_name] = None
                
        # 重新排列列的顺序，年份从新到旧排列
        base_cols = ['school_id', 'school_name', 'spname', 'zs_type', 'sg_name', 'local_batch_name']
        rank_cols = [f'rank_{year}' for year in range(2024, 2017, -1)]
        df = df[base_cols + rank_cols]
        
        return df
    else:
        # 返回空的DataFrame但包含正确的列结构
        base_cols = ['school_id', 'school_name', 'spname', 'zs_type', 'sg_name', 'local_batch_name']
        rank_cols = [f'rank_{year}' for year in range(2024, 2017, -1)]
        return pd.DataFrame(columns=base_cols + rank_cols)

def build_history_rank():
    """构建历史排名数据"""
    print("开始构建历史排名数据...")
    
    # 加载专业分数数据
    prof_data = load_professional_score_data()
    if prof_data is None:
        return
    
    print(f"加载了 {len(prof_data)} 条专业分数记录")
    
    # 分离理科和文科数据
    science_data = prof_data[prof_data['subject'] == 'science'].copy()
    arts_data = prof_data[prof_data['subject'] == 'art'].copy()
    
    print(f"理科记录: {len(science_data)} 条")
    print(f"文科记录: {len(arts_data)} 条")
    
    # 构建理科和文科的历史排名表
    science_rank_df = build_subject_rank_df(science_data, 'science')
    arts_rank_df = build_subject_rank_df(arts_data, 'arts')
    
    # 保存到CSV文件
    os.makedirs('./data', exist_ok=True)
    
    science_output_path = './data/science_history_rank.csv'
    arts_output_path = './data/arts_history_rank.csv'
    
    science_rank_df.to_csv(science_output_path, index=False)
    arts_rank_df.to_csv(arts_output_path, index=False)
    
    print(f"\n历史排名数据已保存:")
    print(f"  理科: {science_output_path} ({len(science_rank_df)} 条记录)")
    print(f"  文科: {arts_output_path} ({len(arts_rank_df)} 条记录)")
    
    # 生成Excel表格
    excel_file = './release/分专业历年录取名次趋势表（2018年~2024年）.xlsx'
    title = '2018年至2024年各专业录取名次趋势表'
    build_excel_table(science_rank_df, arts_rank_df, title, excel_file)
    print(f"\nExcel表格已生成: {excel_file}")

def main():
    build_history_rank()

if __name__ == '__main__':
    main()
    main()
